import numpy as np
from cyw_devkit.core.visualizer import pad_ones, tranform_pt
from cyw_devkit.core.zmath import rotationMatrixToEulerAngles
from copy import deepcopy


class Header:
    def __init__(self, stamp: int, frame_id: str):
        self.stamp = stamp
        self.frame_id = frame_id


class Data:
    def __init__(self, stamp, frame_id, data):
        self.header = Header(stamp, frame_id)
        self.data = data


# 定义边界框角点生成的函数
def get_bbox_corners(center, size, heading_angle):
    # 计算旋转矩阵
    R = np.array([
        [np.cos(heading_angle), -np.sin(heading_angle), 0],
        [np.sin(heading_angle), np.cos(heading_angle), 0],
        [0, 0, 1]
    ])

    # 定义八个角点的偏移量
    half_size = size / 2.0
    corners_offset = np.array([
        [-half_size[0], -half_size[1], -half_size[2]],
        [half_size[0], -half_size[1], -half_size[2]],
        [half_size[0], half_size[1], -half_size[2]],
        [-half_size[0], half_size[1], -half_size[2]],
        [-half_size[0], -half_size[1], half_size[2]],
        [half_size[0], -half_size[1], half_size[2]],
        [half_size[0], half_size[1], half_size[2]],
        [-half_size[0], half_size[1], half_size[2]]
    ])

    # 旋转角点偏移量
    corners_3d = (R @ corners_offset.T).T

    # 加上中心点坐标
    corners_3d += center
    return corners_3d


# 计算每个点是否在边界框内（通过构造AABB包围盒快速检查）
def is_point_in_aabb(point, min_corner, max_corner):
    in_box = (point[None, :, :] >= min_corner[:, None, :]) & (point[None, :, :] <= max_corner[:, None, :])
    in_box = np.all(in_box, axis=-1)
    return in_box


class BBox:
    def __init__(self, data):
        if not isinstance(data, dict):
            raise TypeError("Expected a dictionary for BBox data")
        for key, value in data.items():
            setattr(self, key, value)


class BBoxes:
    def __init__(self, labels):
        if not isinstance(labels, list):
            raise TypeError("Expected a list of dictionaries for BBoxes labels")
        self.obj_ids = []
        self.obj_types = []
        self.tensors = []
        for lab in labels:
            self.obj_ids.append(lab["obj_id"])
            self.obj_types.append(lab["obj_type"])
            boxes = [lab['psr']['position']['x'], lab['psr']['position']['y'], lab['psr']['position']['z'],
                     lab['psr']['scale']['x'], lab['psr']['scale']['y'], lab['psr']['scale']['z'],
                     lab['psr']['rotation']['z']]
            if 'velocity' in lab['psr']:
                boxes.append(lab['psr']['velocity']['x'])
                boxes.append(lab['psr']['velocity']['y'])
                boxes.append(lab['psr']['velocity']['z'])
            self.tensors.append(np.array(boxes))
        self.tensors = np.array(self.tensors)

    def trans_labels(self):
        labels = []
        for obj_id, obj_type, tensor in zip(self.obj_ids, self.obj_types, self.tensors):
            label = dict()
            label['obj_id'] = obj_id
            label['obj_type'] = obj_type
            label['psr'] = dict()
            label['psr']['position'] = dict()
            label['psr']['position']['x'] = tensor[0]
            label['psr']['position']['y'] = tensor[1]
            label['psr']['position']['z'] = tensor[2]
            label['psr']['rotation'] = dict()
            label['psr']['rotation']['x'] = 0
            label['psr']['rotation']['y'] = 0
            label['psr']['rotation']['z'] = tensor[6]
            label['psr']['scale'] = dict()
            label['psr']['scale']['x'] = tensor[3]
            label['psr']['scale']['y'] = tensor[4]
            label['psr']['scale']['z'] = tensor[5]
            labels.append(label)
        return labels

    def filter_by_points(self, lidar):
        if len(self.tensors) == 0:
            return
        # 计算所有边界框的角点
        bbox_corners = np.array([get_bbox_corners(bbox[:3], bbox[3:6], bbox[6]) for bbox in self.tensors])
        # 扩展边界框角点，方便进行点云内点判断
        min_corners = bbox_corners.min(axis=1)  # [10, 3]
        max_corners = bbox_corners.max(axis=1)  # [10, 3]

        pading = np.array([[0.2,0.2,0.0]])
        min_corners-=pading
        max_corners+=pading

        # 使用numpy的向量化比较来判断点是否在AABB包围盒内
        point_mask = is_point_in_aabb(lidar[:, :3], min_corners, max_corners)  # [10000, 10, 1]

        # 计数每个边界框内的点数
        point_counts = np.sum(point_mask, axis=1)

        obj_ids = []
        obj_types = []
        tensors = []
        for obj_id, obj_type, tensor, point_count in zip(self.obj_ids, self.obj_types, self.tensors, point_counts):
            if point_count < 5:
                continue  #
            if obj_type == 'Pedestrian' and point_count < 20:
                continue  # 人box内的点需要大于10
            if obj_type == 'Car' and point_count < 30:
                continue  # box内的点需要大于20
            if (obj_type == 'Truck' or obj_type == 'Bus') and point_count < 40:
                continue  # box内的点需要大于20
            if (obj_type == 'Van') and point_count < 30:
                continue  # box内的点需要大于20
            if (obj_type == 'Bicycle' or obj_type == 'Motorcycle') and point_count < 25:
                continue  # box内的点需要大于15

            obj_ids.append(obj_id)
            obj_types.append(obj_type)
            tensors.append(tensor)
        self.obj_ids = obj_ids
        self.obj_types = obj_types
        self.tensors = np.array(tensors)


def transform_bboxes(bboxes, tf):
    new_boxes = deepcopy(bboxes)
    points3 = bboxes.tensors[:, :3]
    points4 = pad_ones(points3)
    new_boxes.tensors[:, :3] = tranform_pt(points4, tf)[:, :3]
    angle = rotationMatrixToEulerAngles(tf[0][:3, :3])[2]
    new_boxes.tensors[:, 6] = new_boxes.tensors[:, 6] + angle
    return new_boxes
